# -*- coding: utf-8 -*-
"""
211高校一流学科数据爬取与处理脚本
处理rowspan合并行的版本
"""

import requests
from bs4 import BeautifulSoup
import csv
import openpyxl
import re


def crawl_211_universities():
    """
    爬取211高校一流学科数据，正确处理rowspan合并行
    """
    # 设置请求头
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                      "AppleWebKit/537.36 (KHTML, like Gecko) "
                      "Chrome/122.0.0.0 Safari/537.36"
    }

    url = "./index.html"

    try:
        print("正在获取数据...")
        resp = requests.get(url, headers=headers, timeout=15)
        resp.encoding = resp.apparent_encoding
        html = resp.text

        print("数据获取成功，开始解析...")
        soup = BeautifulSoup(html, "html.parser")

        # 查找表格
        table = soup.find("table")
        if not table:
            print("未找到表格")
            return []

        data = []
        rows = table.find_all("tr")

        current_region = None
        region_rowspan = 0

        # 遍历所有行
        for tr in rows:
            tds = tr.find_all("td")
            if not tds:
                continue

            # 跳过表头
            if len(tds) > 0 and ("地区" in tds[0].get_text() or "学校名称" in tds[0].get_text()):
                continue

            # 检查是否有新的地区信息（第一列）
            first_td = tds[0]
            rowspan_attr = first_td.get('rowspan')

            # 如果当前单元格有rowspan属性，说明是新的地区
            if rowspan_attr and rowspan_attr.isdigit():
                region_rowspan = int(rowspan_attr)
                region = first_td.get_text(strip=True)
                if region:
                    current_region = region

                # 处理当前行的学校信息
                if len(tds) >= 2:
                    school = tds[1].get_text(strip=True)
                    subjects = tds[2].get_text(strip=True) if len(tds) >= 3 else ""

                    if current_region and school:
                        data.append([current_region, school, subjects])

                # 减少rowspan计数
                region_rowspan -= 1

            elif region_rowspan > 0 and current_region:
                # 使用当前地区信息
                school = tds[0].get_text(strip=True)
                subjects = tds[1].get_text(strip=True) if len(tds) >= 2 else ""

                if school:
                    data.append([current_region, school, subjects])

                # 减少rowspan计数
                region_rowspan -= 1

            elif len(tds) >= 3:
                # 没有rowspan的情况，每行都有地区信息
                region = tds[0].get_text(strip=True)
                school = tds[1].get_text(strip=True)
                subjects = tds[2].get_text(strip=True)

                if region and school:
                    current_region = region
                    data.append([current_region, school, subjects])

            elif len(tds) >= 2 and current_region:
                # 只有学校和学科，使用当前地区
                school = tds[0].get_text(strip=True)
                subjects = tds[1].get_text(strip=True)

                if school:
                    data.append([current_region, school, subjects])

        print(f"解析完成，共获取 {len(data)} 条记录")
        return data

    except Exception as e:
        print(f"爬取过程中出现错误: {e}")
        import traceback
        traceback.print_exc()
        return []


def save_to_csv(data, filename="211_universities.csv"):
    """
    保存数据到CSV文件
    """
    try:
        with open(filename, "w", newline="", encoding="utf-8-sig") as f:
            writer = csv.writer(f)
            writer.writerow(["地区", "学校名称", "一流学科建设名单"])
            writer.writerows(data)
        print(f"✅ CSV文件已保存: {filename}")
    except Exception as e:
        print(f"保存CSV文件时出现错误: {e}")


def save_to_excel(data, filename="211_universities.xlsx"):
    """
    保存数据到Excel文件
    """
    try:
        wb = openpyxl.Workbook()
        ws = wb.active
        ws.title = "211高校一流学科"

        # 写入表头
        ws.append(["地区", "学校名称", "一流学科建设名单"])

        # 写入数据
        for row in data:
            ws.append(row)

        # 调整列宽
        ws.column_dimensions['A'].width = 15
        ws.column_dimensions['B'].width = 25
        ws.column_dimensions['C'].width = 80

        wb.save(filename)
        print(f"✅ Excel文件已保存: {filename}")
    except Exception as e:
        print(f"保存Excel文件时出现错误: {e}")


def display_statistics(data):
    """
    显示数据统计信息
    """
    if not data:
        print("没有数据可显示")
        return

    # 按地区统计
    region_count = {}
    for row in data:
        region = row[0]
        if region in region_count:
            region_count[region] += 1
        else:
            region_count[region] = 1

    print("\n📊 数据统计:")
    print(f"总计: {len(data)} 所学校")
    print("\n各地区学校数量:")
    for region, count in sorted(region_count.items()):
        region_name = re.sub(r'\(.*?\)', '', region)
        print(f"  {region_name}: {count} 所")


def main():
    """
    主函数
    """
    print("🎓 211高校一流学科数据爬取工具")
    print("=" * 50)

    # 爬取数据
    data = crawl_211_universities()

    if not data:
        print("❌ 未能获取有效数据")
        return

    # 显示统计信息
    display_statistics(data)

    # 保存数据
    save_to_csv(data)
    save_to_excel(data)

    # 显示前几条数据示例
    print("\n📋 数据示例:")
    print("-" * 50)
    for i, (region, school, subjects) in enumerate(data[:5]):
        print(f"{i + 1}. {region} - {school}")
        if len(subjects) > 100:
            print(f"   {subjects[:100]}...")
        else:
            print(f"   {subjects}")
        print()


if __name__ == "__main__":
    main()
